/*******************************************************************************
* Copyright (c) 2008, Jeferson Jose de Miranda                                 *
* UDESC - Universidade do Estado de Santa Catarina                             *
* All rights reserved.                                                         *
*                                                                              *
* Redistribution and use in source and binary forms, with or without           *
* modification, are permitted provided that the following conditions are met:  *
*     * Redistributions of source code must retain the above copyright         *
*       notice, this list of conditions and the following disclaimer.          *
*     * Redistributions in binary form must reproduce the above copyright      *
*       notice, this list of conditions and the following disclaimer in the    *
*       documentation and/or other materials provided with the distribution.   *
*     * Neither the name of the Universidade do Estado de Santa Catarina nor   *
*       the names of its contributors may be used to endorse or promote        *
*       products derived from this software without specific prior written     *
*       permission.                                                            *
*                                                                              *
* THIS SOFTWARE IS PROVIDED BY JEFERSON JOSE DE MIRANDA "AS IS" AND ANY        *
* EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED    *
* WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE       *
* DISCLAIMED. IN NO EVENT SHALL JEFERSON JOSE DE MIRANDA BE LIABLE FOR ANY     *
* DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES   *
* (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; *
* LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND  *
* ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT   *
* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF     *
* THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.            *
*******************************************************************************/

/**
 * \file ImageDifference.hpp
 *
 * \brief Calculates optical flow from the differences between two frames.
 *
 * \author Jeferson Jose de Miranda
 */

#ifndef FLOW_HPP
#define FLOW_HPP

#include "RGBImage.hpp"
#include "GreyLevelImage.hpp"
#include "FlowVectors.hpp"

namespace MoveIn
{
    /**
    * \class Flow
    *
    * \brief Calculates optical flow.
    */
    class Flow
    {
        private:
        Image *current_grey;
        Image *previous_grey;
        Image *current_pyramid;
        Image *previous_pyramid;
        Image *eigen;
        Image *temporary;

        unsigned short int number_of_features;
        unsigned short int width;
        unsigned short int height;
        char *optical_flow_found_feature;
        CvPoint2D32f *current_feature_points;
        CvPoint2D32f *previous_feature_points;
        double quality;
        double minimum_distance;
        int number_of_found_features;
        int window_size;

        int translation_x;
        int translation_y;
        int origin_x;
        int origin_y;

        FlowVectors flow_vectors;

        void alocateFeatureData();

        public:

        Flow(unsigned short int width = 320,
             unsigned short int height = 240,
             unsigned short int number_of_features = 0);

        ~Flow();

        void findGoodFeatures(RGBImage *previous,
                              CvRect *ROI = NULL);

        void setGoodFeatures(PointSequence *features,
                             CvRect *ROI = NULL);

        void updateGoodFeatures();

        void process(RGBImage *previous,
                     RGBImage *current,
                     CvRect *ROI = NULL);

        void process(RGBImage *current,
                     CvRect *ROI = NULL);

        void print(RGBImage *source);

        void printResultVector(RGBImage *source);

        void printGoodFeatures(RGBImage *source);

        void setFlowVectors();

        FlowVectors* getFlowVectors();
    };

/*******************************************************************************
* MoveIn::Flow::Flow                                                           *
*******************************************************************************/

    Flow::Flow(unsigned short int width,
               unsigned short int height,
               unsigned short int number_of_features)
    {
        // Set input image width and height.
        this->width = width;
        this->height = height;
        // Set the number of maximum feature points to track.
        // This values will be used in the method Flow::findGoodFeatures to
        // limit the number of tracking points.
        // Alternatively, a list of points can be set by user with the method
        // Flow::setGoodFeatures. In this case, this value will be reset by a
        // call of this method.
        this->number_of_features = number_of_features;

        // This value is set because of the aperture problem. For more details
        // about this, you will have to search by yourself, sorry.
        window_size = 3;

        // Alocate memory for feature data,
        // only if number of features to track is greater than zero.
        current_feature_points = NULL;
        previous_feature_points = NULL;
        optical_flow_found_feature = NULL;
        if(number_of_features > 0)
        {
            alocateFeatureData();
        }
        // For now, the number of found features is zero.
        number_of_found_features = 0;
        // Allocate all the image/frame buffers.
        // Some of them require single channel float images (this is
        // for workspace on the Lucas an Kanade optical flow algorithm).
        try
        {
            current_grey = new Image(width, height, IPL_DEPTH_8U, 1);
            current_grey->reset();
            previous_grey = new Image(width, height, IPL_DEPTH_8U, 1);
            previous_grey->reset();
            current_pyramid = new Image(width, height, IPL_DEPTH_8U, 1);
            current_pyramid->reset();
            previous_pyramid = new Image(width, height, IPL_DEPTH_8U, 1);
            previous_pyramid->reset();
            eigen = new Image(width, height, IPL_DEPTH_32F, 1);
            eigen->reset();
            temporary = new Image(width, height, IPL_DEPTH_32F, 1);
            temporary->reset();
        }
        catch(std::exception *e)
        {
            throw e;
        }
    }

/*******************************************************************************
* MoveIn::Flow::~Flow                                                          *
*******************************************************************************/

    Flow::~Flow()
    {
        delete current_grey;
        delete previous_grey;
        delete current_pyramid;
        delete previous_pyramid;
        delete eigen;
        delete temporary;
        if(current_feature_points != NULL)
        {
            free(&current_feature_points);
        }
        if(previous_feature_points != NULL)
        {
            free(&previous_feature_points);
        }
        if(optical_flow_found_feature != NULL)
        {
            free(&optical_flow_found_feature);
        }
    }

/*******************************************************************************
* MoveIn::Flow::alocateFeatureData                                             *
*******************************************************************************/

    void Flow::alocateFeatureData()
    {
        // This is for the "ending" points of optical flow vectors.
        if(current_feature_points != NULL)
        {
            free(&current_feature_points);
        }
        current_feature_points =
        (CvPoint2D32f*) malloc(number_of_features*sizeof(CvPoint2D32f));
        if(current_feature_points == NULL)
        {
            throw new OutOfMemory;
        }
        // This is for the "starting" points of optical flow vectors.
        if(previous_feature_points != NULL)
        {
            free(&previous_feature_points);
        }
        previous_feature_points =
        (CvPoint2D32f*) malloc(number_of_features*sizeof(CvPoint2D32f));
        if(previous_feature_points == NULL)
        {
            throw new OutOfMemory;
        }
        // This is for status purposes. The ith position of this array will
        // tell (value other than zero) if the ith feature point could be
        // found in the last analyzed image.
        if(optical_flow_found_feature != NULL)
        {
            free(&optical_flow_found_feature);
        }
        optical_flow_found_feature =
        (char*) malloc(number_of_features*sizeof(char));
        if(optical_flow_found_feature == NULL)
        {
            throw new OutOfMemory;
        }
        // Set status to "none found".
        for(int i=0; i<number_of_features; i++)
        {
            optical_flow_found_feature[i] = 0;
        }
    }

/*******************************************************************************
* MoveIn::Flow::findGoodFeatures                                               *
*******************************************************************************/

    void Flow::findGoodFeatures(RGBImage *previous, CvRect *ROI)
    {
        if(number_of_features == 0)
        {
            return;
        }
        // Set image region of interest, if provided.
        if(ROI != NULL)
        {
            try
            {
                previous->setROI(ROI);
                previous_grey->setROI(ROI);
                eigen->setROI(ROI);
                temporary->setROI(ROI);
            }
            catch(std::exception *e)
            {
                throw e;
            }
        }
        // For now the next two values are hardcoded.
        // Quality is for the minimum quality of the features.
        // Minimum distance specifies the minimum distance between features.
        quality = 0.01;
        minimum_distance = 0.01;
        // Convert reference image to grey scale.
        cvCvtColor(previous->getIplImage(),
                   previous_grey->getIplImage(),
                   CV_BGR2GRAY);
        // Set is variable to number of maximum features to be traceked.
        // The next function will reset its value to the number of
        // found features.
        number_of_found_features = number_of_features;
        // Run Shi and Tomasi good features search algorithm.
        cvGoodFeaturesToTrack(previous_grey->getIplImage(),
                              eigen->getIplImage(),
                              temporary->getIplImage(),
                              previous_feature_points,
                              &number_of_found_features,
                              quality,
                              minimum_distance,
                              NULL);
        // According to OpenCV reference:
        // The function cvFindCornerSubPix iterates to find the sub-pixel
        // accurate location of corners, or radial saddle points.
        cvFindCornerSubPix(current_grey->getIplImage(),
                           current_feature_points,
                           number_of_found_features,
                           cvSize(window_size, window_size),
                           cvSize(-1,-1),
                           cvTermCriteria(CV_TERMCRIT_ITER|CV_TERMCRIT_EPS,
                                          20,
                                          0.03));
        // Reset ROI.
        if(ROI != NULL)
        {
            previous->resetROI();
            previous_grey->resetROI();
            eigen->resetROI();
            temporary->resetROI();
        }
    }

/*******************************************************************************
* MoveIn::Flow::setGoodFeatures                                                *
*******************************************************************************/

    void Flow::setGoodFeatures(PointSequence *features,
                               CvRect *ROI)
    {
        if(features->size() == 0)
        {
            return;
        }
        number_of_features = features->size();
        number_of_found_features = features->size();
        alocateFeatureData();
        optical_flow_found_feature[0] = 1;
        for(int i=0; i<number_of_features; i++)
        {
            optical_flow_found_feature[i] = 1;
            if(ROI == NULL)
            {
                previous_feature_points[i]
                =
                cvPointTo32f(features->getPoint(i));
            }
            else
            {
                previous_feature_points[i]
                =
                cvPointTo32f(cvPoint(features->getPoint(i).x - ROI->x,
                                     features->getPoint(i).y - ROI->y));
            }
        }
    }

/*******************************************************************************
* MoveIn::Flow::updateGoodFeatures                                             *
*******************************************************************************/

    void Flow::updateGoodFeatures()
    {
        // Update the reference values for future interactions.
        previous_grey->copyImage(current_grey);
        previous_pyramid->copyImage(current_pyramid);
        for(int i=0; i<number_of_features; i++)
        {
            previous_feature_points[i]
            =
            current_feature_points[i];
        }
    }

/*******************************************************************************
* MoveIn::Flow::process                                                        *
*******************************************************************************/

    void Flow::process(RGBImage *previous,
                       RGBImage *current,
                       CvRect *ROI)
    {
        // Translation values are for correct drawing the vectors in their
        // positions. Most for debugging purposes.
        translation_x = 0;
        translation_y = 0;
        // Set image ROI if it is not NULL.
        if(ROI != NULL)
        {
            try
            {
                current->setROI(ROI);
                previous->setROI(ROI);
                current_grey->setROI(ROI);
                previous_grey->setROI(ROI);
                current_pyramid->setROI(ROI);
                previous_pyramid->setROI(ROI);
                eigen->setROI(ROI);
                temporary->setROI(ROI);
                // Set translation values according to ROI.
                translation_x = ROI->x;
                translation_y = ROI->y;
                origin_x = translation_x + ROI->width/2;
                origin_y = translation_y + ROI->height/2;
            }
            catch(std::exception *e)
            {
                throw e;
            }
        }
        // Look for good features in previous image.
        findGoodFeatures(previous);

        // Convert current image to grey scale.
        cvCvtColor(current->getIplImage(),
                   current_grey->getIplImage(),
                   CV_BGR2GRAY);
        // If number of found features was set either by
        // Flow::findGoodFeatures or Flow::setGoodFeatures,
        // run Lucas and Kanade optical flow algorithm.
        if(number_of_found_features > 0)
        {
            cvCalcOpticalFlowPyrLK(previous_grey->getIplImage(),
                                   current_grey->getIplImage(),
                                   previous_pyramid->getIplImage(),
                                   current_pyramid->getIplImage(),
                                   previous_feature_points,
                                   current_feature_points,
                                   number_of_features,
                                   cvSize(window_size, window_size),
                                   3,
                                   optical_flow_found_feature,
                                   0,
                                   cvTermCriteria(CV_TERMCRIT_ITER
                                                  |
                                                  CV_TERMCRIT_EPS,
                                                  20,
                                                  0.3),
                                   0);//flags);

            //flags |= CV_LKFLOW_PYR_A_READY;
        }
        // Reset ROI.
        if(ROI != NULL)
        {
            current->resetROI();
            previous->resetROI();
            current_grey->resetROI();
            previous_grey->resetROI();
            current_pyramid->resetROI();
            previous_pyramid->resetROI();
            eigen->resetROI();
            temporary->resetROI();
        }
        // Set the results in the proper data structure.
        setFlowVectors();
    }

/*******************************************************************************
* MoveIn::Flow::process                                                        *
*******************************************************************************/

    void Flow::process(RGBImage *current,
                       CvRect *ROI)
    {
        if(number_of_found_features == 0)
        {
            return;
        }
        // Translation values are for correct drawing the vectors in their
        // positions. Most for debugging purposes.
        translation_x = 0;
        translation_y = 0;
        // Set image ROI if it is not NULL.
        if(ROI != NULL)
        {
            try
            {
                current->setROI(ROI);
                current_grey->setROI(ROI);
                previous_grey->setROI(ROI);
                current_pyramid->setROI(ROI);
                previous_pyramid->setROI(ROI);
                eigen->setROI(ROI);
                temporary->setROI(ROI);
                // Set translation values according to ROI.
                translation_x = ROI->x;
                translation_y = ROI->y;
                origin_x = translation_x + ROI->width/2;
                origin_y = translation_y + ROI->height/2;
            }
            catch(std::exception *e)
            {
                throw e;
            }
        }
        // Convert current image to grey scale.
        cvCvtColor(current->getIplImage(),
                   current_grey->getIplImage(),
                   CV_BGR2GRAY);
        // If number of found features was set either by
        // Flow::findGoodFeatures or Flow::setGoodFeatures,
        // run Lucas and Kanade optical flow algorithm.
        cvCalcOpticalFlowPyrLK(previous_grey->getIplImage(),
                               current_grey->getIplImage(),
                               previous_pyramid->getIplImage(),
                               current_pyramid->getIplImage(),
                               previous_feature_points,
                               current_feature_points,
                               number_of_features,
                               cvSize(window_size, window_size),
                               3,
                               optical_flow_found_feature,
                               0,
                               cvTermCriteria(CV_TERMCRIT_ITER
                                              |
                                              CV_TERMCRIT_EPS,
                                              20,
                                              0.3),
                               0);//flags);
        //flags |= CV_LKFLOW_PYR_A_READY;

        // Reset ROI.
        if(ROI != NULL)
        {
            current->resetROI();
            current_grey->resetROI();
            previous_grey->resetROI();
            current_pyramid->resetROI();
            previous_pyramid->resetROI();
            eigen->resetROI();
            temporary->resetROI();
        }
        //updateGoodFeatures();
        // Set the results in the proper data structure.
        setFlowVectors();
    }

/*******************************************************************************
* MoveIn::Flow::print                                                          *
*******************************************************************************/

    void Flow::print(RGBImage *source)
    {
        flow_vectors.print(source, 10.0);
    }

/*******************************************************************************
* MoveIn::Flow::printResultVector                                              *
*******************************************************************************/

    void Flow::printResultVector(RGBImage *source)
    {
        //flow_vectors.sum();
        flow_vectors.result();
        flow_vectors.print(source, 10.0);
    }

/*******************************************************************************
* MoveIn::Flow::printGoodFeatures                                              *
*******************************************************************************/

    void Flow::printGoodFeatures(RGBImage *source)
    {
        if(number_of_found_features == 0)
        {
            return;
        }
        for(int i=0; i<number_of_features; i++)
        {
            if(optical_flow_found_feature[i] == 0)
            {
                continue;
            }
            CvPoint good_feature;
            good_feature.x = (int) previous_feature_points[i].x;
            good_feature.y = (int) previous_feature_points[i].y;
            if(translation_x != 0 || translation_y != 0)
            {
                good_feature.x += translation_x;
                good_feature.y += translation_y;
            }
            cvCircle(source->getIplImage(),
                     good_feature,
                     3,
                     CV_RGB(255,255,000),
                     4);
        }
    }

/*******************************************************************************
* MoveIn::Flow::setFlowVectors                                                 *
*******************************************************************************/

    void Flow::setFlowVectors()
    {
        flow_vectors.reset();
        for(int i=0; i<number_of_features; i++)
        {
            if(optical_flow_found_feature[i] == 0)
            {
                continue;
            }
            CvPoint p, q;
            p.x = (int) previous_feature_points[i].x;
            p.y = (int) previous_feature_points[i].y;
            q.x = (int) current_feature_points[i].x;
            q.y = (int) current_feature_points[i].y;
            flow_vectors.append(p, q);
        }
        if(translation_x != 0 || translation_y != 0)
        {
            flow_vectors.translate(translation_x, translation_y);
            flow_vectors.setOrigin(cvPoint(origin_x, origin_y));
        }
    }

/*******************************************************************************
* MoveIn::Flow::getFlowVectors                                                 *
*******************************************************************************/

    FlowVectors* Flow::getFlowVectors()
    {
        return &flow_vectors;
    }
}

#endif
